import re
import json
from pathlib import Path
from collections import deque
from datetime import datetime
from typing import TYPE_CHECKING, Dict

from mozi_ai_x.utils import mprint_with_name
from mozi_ai_x.simulation.active_unit.aircraft import CAircraft
from mozi_ai_x.simulation.active_unit.ship import CShip
from mozi_ai_x.simulation.active_unit.submarine import CSubmarine
from vectorvein.types.enums import BackendType
from vectorvein.chat_clients import create_async_chat_client

from heli_anti_sub.utils.docstring_parser import parse_class_methods

if TYPE_CHECKING:
    from mozi_ai_x.simulation.side import CSide
    from mozi_ai_x.simulation.scenario import CScenario
    from .env import Environment

mprint = mprint_with_name(__name__)


class ConversationHistory:
    def __init__(self, max_length=10, save_folder: str | Path | None = None):
        self.history = deque(maxlen=max_length)
        self.save_folder = Path(save_folder) if save_folder else Path.cwd()
        self.create_time = datetime.now().strftime("%Y%m%d_%H%M%S")

    def add_message(self, role: str, content: str):
        if len(self.history) > 0 and self.history[-1]["role"] == role:
            mprint.warning(f"检测到连续的 {role} 消息，移除上一条消息")
            self.history.pop()
        self.history.append({"role": role, "content": content})

    def get_messages(self):
        return list(self.history)

    def save_messages(self):
        """将对话历史保存为文件"""

        messages = self.get_messages()
        with open(self.save_folder / f"{self.create_time}.md", "w", encoding="utf-8") as f:
            for message in messages:
                f.write(f"===={message['role']}====\n")
                if "reasoning_content" in message:
                    f.write("-------reasoning_content-------\n")
                    f.write(f"{message['reasoning_content']}\n")
                f.write("-------content-------\n")
                f.write(f"{message['content']}\n\n")

        with open("conversations.json", "w", encoding="utf8") as f:
            json.dump(messages, f, ensure_ascii=False, indent=2)


class Agent:
    def __init__(self, env: "Environment"):
        super(Agent, self).__init__()
        self.env = env

        # 初始化 LLM 相关组件
        self.chat_client = create_async_chat_client(backend=BackendType.DeepSeek, model="deepseek-reasoner")
        self.conversation_history = ConversationHistory(max_length=10, save_folder="log")

        # 初始化单元类型到可用工具的映射
        self.unit_tools = {}
        self._init_unit_tools()

    def _init_unit_tools(self):
        """初始化不同类型单元可用的工具"""

        # 解析各类单元的方法
        self.unit_tools = {
            "aircraft": parse_class_methods(CAircraft),
            "ship": parse_class_methods(CShip),
            "submarine": parse_class_methods(CSubmarine),
        }

    def _generate_tools_description(self, unit_type: str) -> str:
        """生成指定单元类型的工具描述

        Args:
            unit_type: 单元类型 (aircraft/ship/submarine)

        Returns:
            str: 工具描述文本
        """
        if unit_type not in self.unit_tools:
            return ""

        tools_text = json.dumps(self.unit_tools[unit_type], ensure_ascii=False)
        return tools_text

    def _group_units_by_type(self, side: "CSide") -> Dict[str, list]:
        """按类型对单元进行分组

        Args:
            side: 推演方对象

        Returns:
            Dict[str, list]: 按类型分组的单元列表
        """
        return {
            "aircraft": list(side.aircrafts.values()),
            "ship": list(side.ships.values()),
            "submarine": list(side.submarines.values()),
        }

    @staticmethod
    def _construct_active_unit_text(unit: "CAircraft | CShip | CSubmarine", indent: int = 4) -> str:
        """构建单元信息文本

        Args:
            unit: 单元对象

        Returns:
            str: 单元信息文本
        """

        indent_prefix = " " * indent

        text = f"{indent_prefix}- {unit.name}({unit.guid}):\n"
        text += f"{indent_prefix}  位置: 经度 {unit.longitude}, 纬度 {unit.latitude}, 海拔 {unit.altitude}\n"
        text += f"{indent_prefix}  朝向: {unit.current_heading}, 速度: {unit.current_speed}\n"
        text += f"{indent_prefix}  油门: {unit.current_throttle}\n"
        text += f"{indent_prefix}  状态: {unit.active_unit_status.replace('状态:', '').strip()}\n"
        text += f"{indent_prefix}  受损: {unit.damage_state}\n"

        sensors = unit.get_sensor()
        if sensors:
            text += f"{indent_prefix}  传感器:\n"
            for sensor in sensors.values():
                text += f"{indent_prefix}    {sensor.name}({sensor.guid}):\n"
                text += f"{indent_prefix}      类型: {sensor.sensor_type_label}\n"
                text += f"{indent_prefix}      状态: {sensor.component_status_label}\n"
                text += f"{indent_prefix}      毁伤: {sensor.damage_severity_label}\n"

        text += f"{indent_prefix}  燃油状态:\n"
        fule_state_lines = unit.fuel_state.split("\n")
        for line in fule_state_lines:
            text += f"{indent_prefix}      {line.strip()}\n"

        weapons = unit.get_weapon_infos()
        if weapons:
            text += f"{indent_prefix}  武器状态:\n"
            for weapon in weapons:
                text += f"{indent_prefix}    {weapon}\n"

        text += f"{indent_prefix}  任务: {unit.get_assigned_mission()}\n"

        return text

    @staticmethod
    def extract_json_from_text(text: str):
        """从文本中提取JSON字符串"""
        json_str = re.search(r"```json(.*?)```", text, re.DOTALL)
        if json_str:
            return json.loads(json_str.group(1))
        return None

    @staticmethod
    def get_battlefield_state_text(side: "CSide", unit_type: str | None = None) -> str:
        """构建战场信息文本

        Args:
            side: 推演方对象
            unit_type: 可选，指定要描述的单元类型

        Returns:
            str: 战场态势文本
        """
        text = ""

        # 根据指定类型输出单元信息
        if unit_type == "aircraft" or unit_type is None:
            if side.aircrafts:
                text += "飞机：\n"
                for aircraft in side.aircrafts.values():
                    text += Agent._construct_active_unit_text(aircraft, 4)

        if unit_type == "ship" or unit_type is None:
            if side.ships:
                text += "水面舰艇：\n"
                for ship in side.ships.values():
                    text += Agent._construct_active_unit_text(ship, 4)

        if unit_type == "submarine" or unit_type is None:
            if side.submarines:
                text += "潜艇：\n"
                for submarine in side.submarines.values():
                    text += Agent._construct_active_unit_text(submarine, 4)

        # 探测到的目标信息
        text += "探测到的目标：\n"
        for contact in side.contacts.values():
            text += f"    - {contact.name} ({contact.guid}):\n"
            text += f"      类型: {contact.get_type_description()}\n"
            text += f"      位置: 经度 {contact.longitude}, 纬度 {contact.latitude}\n"

        return text

    async def get_ai_response(self, text: str, unit_type: str, role: str):
        """获取 AI 的响应

        Args:
            text: 战场态势文本
            unit_type: 单元类型
        """

        def get_prompt(prompt_name: str, params: dict):
            prompt_path = Path(__file__).parent.parent / "prompts" / f"{prompt_name}.md"
            with open(prompt_path, "r", encoding="utf-8") as f:
                prompt = f.read()

            # 替换其他参数
            for key, value in params.items():
                prompt = prompt.replace(f"{{{{{key}}}}}", value)

            return prompt

        tools_desc = self._generate_tools_description(unit_type)
        prompt = get_prompt("prompt2", {"role": role, "unit_type": unit_type, "text": text, "tools": tools_desc})
        mprint.info("Prompt:")
        mprint.info(prompt)

        self.conversation_history.add_message("user", prompt)
        messages = self.conversation_history.get_messages()

        reasoning_content = ""
        content = ""
        reasoning_end = False

        response = await self.chat_client.create_stream(messages=messages, temperature=0.6)
        prompt_tokens = completion_tokens = 0
        async for chunk in response:
            if not reasoning_end:
                if chunk.reasoning_content:
                    reasoning_content += chunk.reasoning_content
                    print(chunk.reasoning_content, end="")
            if chunk.content:
                if not reasoning_end:
                    reasoning_end = True
                    print("\n" + "=" * 30)
                content += chunk.content
                print(chunk.content, end="")

            if chunk.usage and chunk.usage.prompt_tokens:
                prompt_tokens = chunk.usage.prompt_tokens
            if chunk.usage and chunk.usage.completion_tokens:
                completion_tokens = chunk.usage.completion_tokens

        mprint.debug(f"Prompt tokens: {prompt_tokens}, Completion tokens: {completion_tokens}")

        json_str = self.extract_json_from_text(content or "")

        return {
            "think": reasoning_content,
            "text": content,
            "actions": json_str["actions"] if json_str else [],
        }

    async def execute_commands(self, scenario: "CScenario", actions: list):
        """执行 LLM 生成的命令"""
        red_side = scenario.get_side_by_name("美国海军")
        if not red_side:
            raise ValueError("美国海军不存在")

        for action in actions:
            unit_guid = action["unit_guid"]
            method_name = action["method"]
            params = action["parameters"]

            mprint.info(f"执行动作: {method_name}")
            mprint.info(f"单元: {unit_guid}")
            mprint.info(f"参数: {params}")

            try:
                # 获取单元实例
                unit = None
                if unit_guid in red_side.aircrafts:
                    unit = red_side.aircrafts[unit_guid]
                elif unit_guid in red_side.ships:
                    unit = red_side.ships[unit_guid]
                elif unit_guid in red_side.submarines:
                    unit = red_side.submarines[unit_guid]

                if unit is None:
                    mprint.warning(f"未找到指定单元: {unit_guid}")
                    continue

                # 获取方法
                method = getattr(unit, method_name, None)
                if method is None:
                    mprint.warning(f"单元 {unit_guid} 不支持方法: {method_name}")
                    continue

                # 调用方法
                result = await method(**params)
                mprint.info(f"执行结果: {result}")

            except Exception as e:
                mprint.error(f"执行 {method_name} 时发生错误: {str(e)}")

    async def step(self):
        """执行一步智能体行动"""
        # 获取当前战场态势
        red_side = self.env.scenario.get_side_by_name("美国海军")
        if red_side is None:
            raise ValueError("美国海军不存在")

        # 按单元类型分组
        units_by_type = self._group_units_by_type(red_side)

        # 对每种类型的单元进行分析和决策
        for unit_type, units in units_by_type.items():
            if not units:  # 跳过没有单元的类型
                continue

            # 获取该类型单元的战场态势
            battlefield_text = self.get_battlefield_state_text(red_side, unit_type)

            # 获取 AI 响应
            ai_response = await self.get_ai_response(battlefield_text, unit_type, red_side.name)
            self.conversation_history.add_message("assistant", ai_response["text"])
            mprint.info(f"{unit_type} AI响应:", ai_response)

            self.conversation_history.save_messages()

            # 执行命令
            json_data = self.extract_json_from_text(ai_response["text"])
            if json_data and "actions" in json_data:
                await self.execute_commands(self.env.scenario, json_data["actions"])

        return
